from datetime import datetime
from fastapi import APIRouter
from pydantic import BaseModel
from src.huey_init.storages import Storage
from src.huey_init._huey import huey
import json

from src.huey_init.task_status import storage
from src.tasks.task_queue import get_tasks_queue, delete_running_key
from src.tasks.task_redis import redis_queue_change
from src.huey_init.pgsql import db_manager
import time
import os
from config import PM2ID


class Priority(BaseModel):
    FunctionType: int
    TaskType: int
    TaskList: list

router = APIRouter(
    prefix="/tasks",
    tags=["tasks"],
    responses={404: {"description": "Not found"}}
)

@router.get("/")
async def tasks():
    results = get_tasks_queue()
    return {"status":"success", "taskInfoList":results}

@router.get("/status/{task_id}")
async def get_task_status(task_id: str):
    res = db_manager.fetch_all("public.task_status", columns="status", condition={"task_id": task_id})
    # res: [('success',)]
    if not res:
        return {"status": "not_executed"}
    else:
        return {"status": res[0][0]}

@router.post("/priority")
async def task_priority_change(tasks_data: Priority):
    try:
        task_list = []
        for task in tasks_data.TaskList:
            ai_task_id = task["AiTaskId"]
            task_list.append(ai_task_id)
        # print("start: ", datetime.now())
        redis_queue_change(task_list)
        # print("end: ", datetime.now())
        return {"status": "success"}
    except Exception as e:
        print(f"Error priority update: {e}")
        return {"status": "error"}

@router.post("/cancel/{task_id}")
async def task_cancel(task_id: str):
    res = await get_task_status(task_id)
    try:
        if res["status"] == "running":  # 取消正在执行的的任务
            pm2_cmd = "pm2 restart {}".format(PM2ID)
            delete_running_key(task_id)
            db_manager.delete("public.task_status", {'task_id': task_id})
            os.system(pm2_cmd)
        else:  # 取消待执行的任务
            huey.revoke_by_id(task_id)
            storage.sets("revoke_tasks:{}".format(task_id), {"id": task_id}, 43200)
        return {"status": "success"}
    except Exception as e:
        print(f"Error canceling task: {e}")
        return {"status": "error"}

